import torch
from data_loader import data_load
from data_iterator import data_iter_random
from model import RNNModel
from torch import nn

def predict(prefix, num_chars, model, idx_to_char, char_to_idx):
    """
    使用训练好的RNN模型进行预测。

    参数:
    prefix (str): 预测的起始字符串。
    num_chars (int): 预测的字符数。
    model (RNNModel): 训练好的RNN模型。
    idx_to_char (list): 索引到字符的映射。
    char_to_idx (dict): 字符到索引的映射。

    返回:
    str: 预测的字符串。
    """
    state = None
    output = []
    # 预热模型，即输入初始字符串
    for char in prefix:
        _, state = model(torch.tensor(char_to_idx[char]).unsqueeze(0).unsqueeze(0), state)
        output.append(char_to_idx[char])
    # 生成预测的字符
    for _ in range(num_chars - len(prefix)):
        X = torch.tensor(output[-1]).unsqueeze(0).unsqueeze(0)
        (Y, state) = model(X, state)
        output.append(int(Y.argmax(dim=1).item()))
    return ''.join([idx_to_char[i] for i in output])

def train_and_predict_rnn(model, texts_idx, num_epochs, batch_size, num_steps, lr, idx_to_char, char_to_idx, prefix,
                          num_chars):
    """
    训练RNN模型并进行预测。

    参数:
    model (RNNModel): RNN模型。
    texts_idx (list): 文本数据的索引表示。
    num_epochs (int): 训练的轮数。
    batch_size (int): 批次大小。
    num_steps (int): 时间步的长度。
    lr (float): 学习率。
    idx_to_char (list): 索引到字符的映射。
    char_to_idx (dict): 字符到索引的映射。
    prefix (str): 预测的起始字符串。
    num_chars (int): 预测的字符数。
    """
    trainer = torch.optim.Adam(model.parameters(), lr=lr)
    loss = nn.CrossEntropyLoss()
    for epoch in range(num_epochs):
        loss_sum = 0
        # 使用随机采样的数据迭代器进行训练
        for X, Y in data_iter_random(texts_idx, batch_size, num_steps):
            state = None
            output, state = model(X, state)
            Y = Y.transpose(0, 1).reshape(-1)
            l = loss(output, Y.long())
            trainer.zero_grad()
            l.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.01)
            trainer.step()
            loss_sum += l.item() * Y.shape[0]
        print(f"epoch: {epoch}, loss: {loss_sum / len(texts_idx)}")
        if epoch % 10 == 0:
            print(predict(prefix, num_chars, model, idx_to_char, char_to_idx))

if __name__ == "__main__":
    file_path =  "./data/jaychou_lyrics.txt.zip"
    texts, idx_to_char, char_to_idx, vocab_size = data_load(10000, file_path)
    texts_idx = [char_to_idx[c] for c in texts]
    model = RNNModel(input_size=vocab_size, num_hiddens=256, num_layers=1)
    train_and_predict_rnn(model, texts_idx, 201, 32, 25, 0.01, idx_to_char, char_to_idx, "晴天", 50)
